package com.xunlai.infra.permission.manager.authority.sercurity.user;

import com.xunlai.infra.permission.common.constant.ExtendClaimsConstants;
import com.xunlai.infra.permission.common.security.TenantUser;
import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
import org.springframework.security.core.GrantedAuthority;
import org.springframework.security.oauth2.core.OAuth2AccessToken;
import org.springframework.security.oauth2.core.oidc.OidcIdToken;
import org.springframework.security.oauth2.core.oidc.OidcScopes;
import org.springframework.security.oauth2.core.oidc.OidcUserInfo;
import org.springframework.security.oauth2.core.oidc.StandardClaimNames;
import org.springframework.security.oauth2.server.authorization.OAuth2Authorization;
import org.springframework.security.oauth2.server.authorization.oidc.authentication.OidcUserInfoAuthenticationContext;

import java.util.*;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * @author liang
 * @date 12/8/22 4:25 PM
 */
public class EnhanceOidcUserInfoMapper implements Function<OidcUserInfoAuthenticationContext, OidcUserInfo> {
    private static final List<String> EMAIL_CLAIMS = Arrays.asList(
            StandardClaimNames.EMAIL,
            StandardClaimNames.EMAIL_VERIFIED
    );
    private static final List<String> PHONE_CLAIMS = Arrays.asList(
            StandardClaimNames.PHONE_NUMBER,
            StandardClaimNames.PHONE_NUMBER_VERIFIED
    );
    private static final List<String> PROFILE_CLAIMS = Arrays.asList(
            StandardClaimNames.NAME,
            StandardClaimNames.FAMILY_NAME,
            StandardClaimNames.GIVEN_NAME,
            StandardClaimNames.MIDDLE_NAME,
            StandardClaimNames.NICKNAME,
            StandardClaimNames.PREFERRED_USERNAME,
            StandardClaimNames.PROFILE,
            StandardClaimNames.PICTURE,
            StandardClaimNames.WEBSITE,
            StandardClaimNames.GENDER,
            StandardClaimNames.BIRTHDATE,
            StandardClaimNames.ZONEINFO,
            StandardClaimNames.LOCALE,
            StandardClaimNames.UPDATED_AT
    );

    private static final String ID_CLAIM = ExtendClaimsConstants.ID;
    private static final String ROLE_CLAIM = ExtendClaimsConstants.ROLES;
    private static final List<String> PERMISSION_CLAIMS = Arrays.asList(
            ExtendClaimsConstants.TENANT,
            "department"
    );

    @Override
    public OidcUserInfo apply(OidcUserInfoAuthenticationContext authenticationContext) {
        OAuth2Authorization authorization = authenticationContext.getAuthorization();
        OidcIdToken idToken = authorization.getToken(OidcIdToken.class).getToken();
        OAuth2AccessToken accessToken = authenticationContext.getAccessToken();
        Map<String, Object> enhancedClaims = new HashMap<>(idToken.getClaims()) ;
        UsernamePasswordAuthenticationToken authenticationToken = authorization.getAttribute("java.security.Principal");
        Collection<GrantedAuthority> authorities = authenticationToken.getAuthorities();
        Set<String> roles = authorities.stream().map(GrantedAuthority::getAuthority).collect(Collectors.toSet());
        enhancedClaims.put(ExtendClaimsConstants.ROLES,roles);
        Object principal = authenticationToken.getPrincipal();
        if(principal instanceof TenantUser tenantUser){
            enhancedClaims.put(ExtendClaimsConstants.TENANT,tenantUser.getTenant());
            enhancedClaims.put(ExtendClaimsConstants.ID,tenantUser.getId());
        }
        Map<String, Object> scopeRequestedClaims = getClaimsRequestedByScope(enhancedClaims,
                accessToken.getScopes());

        return new OidcUserInfo(scopeRequestedClaims);
    }

    private static Map<String, Object> getClaimsRequestedByScope(Map<String, Object> claims, Set<String> requestedScopes) {
        Set<String> scopeRequestedClaimNames = new HashSet<>(32);
        scopeRequestedClaimNames.add(StandardClaimNames.SUB);
        scopeRequestedClaimNames.add(ID_CLAIM);

        if(requestedScopes.contains("all")){
            return claims;
        }

        if (requestedScopes.contains(OidcScopes.ADDRESS)) {
            scopeRequestedClaimNames.add(StandardClaimNames.ADDRESS);
        }
        if (requestedScopes.contains(OidcScopes.EMAIL)) {
            scopeRequestedClaimNames.addAll(EMAIL_CLAIMS);
        }
        if (requestedScopes.contains(OidcScopes.PHONE)) {
            scopeRequestedClaimNames.addAll(PHONE_CLAIMS);
        }
        if (requestedScopes.contains(OidcScopes.PROFILE)) {
            scopeRequestedClaimNames.addAll(PROFILE_CLAIMS);
        }

        if(requestedScopes.contains("role")){
            scopeRequestedClaimNames.add(ROLE_CLAIM);
        }
        if(requestedScopes.contains("permission")){
            scopeRequestedClaimNames.addAll(PERMISSION_CLAIMS);
        }

        Map<String, Object> requestedClaims = new HashMap<>(claims);
        requestedClaims.keySet().removeIf(claimName -> !scopeRequestedClaimNames.contains(claimName));

        return requestedClaims;
    }
}
